import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from model.radiation_transport_model import radiation_transport_model

# 读取CSV文件
environment_data = pd.read_csv('data_resource\environmental_data.csv')
environment_data['timestamp'] = pd.to_datetime(environment_data['timestamp'])

# 读取已有的辐射数据
radiation_data = pd.read_csv('data_resource\\radiation_data.csv')
radiation_data['timestamp'] = pd.to_datetime(radiation_data['timestamp'])

# 辐射源特性
radiation_source_properties = {'source_strength': 100, 'half_life': 10}

# 预测的辐射浓度集
predicted_concentrations = []

# 针对环境数据预测辐射浓度
for idx, row in environment_data.iterrows():
    # 获取当前环境数据
    env_data = environment_data.loc[idx]
    
    # 找到最接近当前环境数据时间的辐射数据
    nearest_time = min(radiation_data['timestamp'], key=lambda x: abs(x - env_data['timestamp']))
    nearest_rad_data = radiation_data[radiation_data['timestamp'] == nearest_time]
    
    # 预测辐射浓度
    predicted_concentration = radiation_transport_model(env_data, radiation_source_properties)
    predicted_concentrations.append(predicted_concentration)

# 创建数据框
data = pd.DataFrame({
    'timestamp': environment_data['timestamp'],
    'concentration': predicted_concentrations
})

# 保存数据到 CSV 文件
data.to_csv('data_resource\\radiation_data2.csv', index=False)
# 绘制预测图表
plt.figure(figsize=(10, 6))
plt.plot(environment_data['timestamp'], predicted_concentrations)
plt.title('Predicted Radiation Concentration over Time')
plt.xlabel('Timestamp')
plt.ylabel('Predicted Concentration (Bq/m^3)')
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()
